package top.dreamcenter.ai.llm.llmchat.config;

import jakarta.annotation.Resource;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.memory.InMemoryChatMemory;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.mcp.SyncMcpToolCallbackProvider;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
public class ChatClientRegister {

    @Resource
    private ChatModel chatModel;

    @Resource
    private SyncMcpToolCallbackProvider toolCallbackProvider;

    @Resource
    private SimpleVectorStore vectorStore;

    @Resource
    private InMemoryChatMemory inMemoryChatMemory;


    /**
     * mcp函数调用
     * @return client
     */
    @Bean("toolCallClient")
    public ChatClient toolCallClient() {
        return ChatClient
                .builder(chatModel)
                .defaultTools(toolCallbackProvider.getToolCallbacks())
                .build();
    }

    @Bean("ragRetrieveClient")
    public ChatClient ragRetrieveClient() {
        return ChatClient
                .builder(chatModel)
                .defaultAdvisors(
                        new QuestionAnswerAdvisor(vectorStore), // RAG
                        new MessageChatMemoryAdvisor(
                                inMemoryChatMemory,
                                AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY,
                                10), // chat memory
                        new SimpleLoggerAdvisor() // log
                )
                .build();
    }

}
